import jwt
import datetime
import uuid
from typing import Tuple, Union, Any, Dict
from . import config
from . import code


class TokenManager:
    def __init__(self, secret_key: str = None, expires: int = None, refresh_expires=None, algorithm="HS256"):
        """
        :param secret_key:  密匙
        :param expires:  过期时间
        :param refresh_expires:  刷新token过期时间
        :param algorithm:
        """
        self.secret_key = secret_key or config.SECRET_KEY
        self.expires = expires or config.EXPIRES
        self.algorithm = algorithm
        self.refresh_expires = refresh_expires or self.expires * 2

    @property
    def access_id(self):
        return str(uuid.uuid4())

    # 生成token
    def gen_jwt(self, **kwargs) -> Tuple[str, str, str, str]:
        """
        :param kwargs:
        :return:
        """
        payload = self._get_access_payload(**kwargs)
        refresh_payload = self._get_refresh_payload(**kwargs)

        # token 过期时间
        access_expires = self.get_access_expires()

        access_token = jwt.encode(payload, self.secret_key, self.algorithm)
        refresh_token = jwt.encode(refresh_payload, self.secret_key, self.algorithm)
        return access_token, refresh_token, str(access_expires), self.access_id

    # 刷新
    def refresh_jwt(self, refresh_token):
        success, payload = self.verify_jwt(refresh_token)
        assert payload["type"] == 'refresh'
        return self.gen_jwt(**payload), payload

    # 验证token
    def verify_jwt(self, access_token: str) -> Union[Tuple[int, dict], Tuple[int, str]]:
        try:
            payload: dict = jwt.decode(access_token, self.secret_key, self.algorithm)
            return code.TOKEN_SUCCESS, payload
        except jwt.exceptions.ExpiredSignatureError:
            return code.TOKEN_OVERDUE, "token已过期"
        except jwt.DecodeError:
            return code.TOKEN_ACCESS, "token认证失败"
        except jwt.InvalidTokenError:
            return code.TOKEN_FAILED, "非法的token"

    def _get_access_payload(self, **kwargs) -> Dict[str, Any]:
        # token过期时间
        exp = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=self.expires)
        kwargs.update({
            "exp": exp,
            "type": "access",
        })
        return kwargs

    def _get_refresh_payload(self, **kwargs) -> Dict[str, Any]:
        """
        :param kwargs:
        :return:   refresh token  payload
        """
        exp = datetime.datetime.now(tz=datetime.timezone.utc) + datetime.timedelta(seconds=self.refresh_expires)
        kwargs.update({
            "exp": exp,
            "type": "refresh",
        })
        return kwargs

    def get_access_expires(self) -> datetime:
        """
        :return: access token expires
        """
        exp = datetime.datetime.now() + datetime.timedelta(seconds=self.expires)
        return exp
